import requests
import json
from langchain.prompts import PromptTemplate
from boat_nl2sql.utiles import extract_sql_from_markdown, load_table_info_from_json
from boat_nl2sql.oracle_connection import get_db_connection, FewShotPromptTemplate

# 在这里导入您的 LLM
from boat_nl2sql.llm_module import SiliconFlow
llm = SiliconFlow()

# API配置
API_URL = 'http://10.1.161.53:11434/api/chat'
MODEL_NAME = "deepseek-r1:70b"

def generate_sql_data(question):

    print("\n===== 测试 SQL 查询链 =====")

    # 加载表信息
    full_table_info = load_table_info_from_json("../boat_nl2sql/oracle_table_info.json")
    include_tables = [
        'BAYONET_BUSSINESS.BOAT_WARNING',
        'BAYONET_BUSSINESS.OFF_SITE_CASE',
        'BAYONET_BUSSINESS.ZD_WARNING_TYPE',
        'BAYONET_DYNAMIC.VIDEO_ANALYSIS',
        'BAYONET_DYNAMIC.DATAFUSION',
        'BAYONET_BASICS.SYS_BAYONET'
    ]

    # 获取数据库连接
    db = get_db_connection(
        sample_rows=2,
        custom_table_info=full_table_info,
        include_tables=include_tables
    )

    # 配置few_shot提示模板
    examples = [
        {
            "input": "查询名称中包含'船闸'的所有卡口基本信息，按名称排序。",
            "query": "SELECT ID, NAME, UNIT_ID, DES, PHOTOHTTP  FROM BAYONET_BASICS.SYS_BAYONET  WHERE NAME LIKE '%船闸%'  ORDER BY NAME",
        },
        {
            "input": "统计一月份各卡口点的船舶通过数量，按照通过船舶数量降序排列，只显示通过量前10的卡口。",
            "query": "SELECT * FROM (SELECT b.NAME AS 卡口名称, COUNT(d.CODE) AS 通过船舶数量  FROM BAYONET_BASICS.SYS_BAYONET b  JOIN BAYONET_DYNAMIC.DATAFUSION d ON b.ID = d.BAYONET_ID  WHERE d.PASSTIME >= TRUNC(SYSDATE, 'YYYY') AND d.PASSTIME < ADD_MONTHS(TRUNC(SYSDATE, 'YYYY'), 1) GROUP BY b.NAME  ORDER BY 通过船舶数量 DESC) WHERE ROWNUM <= 10",
        },
        {
            "input": "查询系统中各类型预警的发生次数和占比，按发生次数从高到低排序。",
            "query": "WITH warning_counts AS (  SELECT t.NAME AS 预警类型,   COUNT(w.CODE) AS 预警次数  FROM BAYONET_BUSSINESS.ZD_WARNING_TYPE t  JOIN BAYONET_BUSSINESS.BOAT_WARNING w ON t.ID = w.WARNING_TYPE  GROUP BY t.NAME  )  SELECT    预警类型,预警次数, ROUND(预警次数 * 100 / SUM(预警次数) OVER(), 2) || '%' AS 占比 FROM warning_counts ORDER BY 预警次数 DESC "
        },
        {
            "input": "分析一下2025年1月油墩港(东大盈船闸)的流量组成。",
            "query": """SELECT 
        COUNT(da.CODE) AS TOTAL_FLOW,          -- 总流量
        SUM(CASE WHEN da.DIRECTION = 1 THEN 1 ELSE 0 END) AS UPSTREAM_FLOW, -- 上行流量
        SUM(CASE WHEN da.DIRECTION = 2 THEN 1 ELSE 0 END) AS DOWNSTREAM_FLOW -- 下行流量
    FROM 
        BAYONET_DYNAMIC.DATAFUSION da
    LEFT JOIN BAYONET_BASICS.SYS_BAYONET sb
        ON da.BAYONET_ID = sb.ID
    WHERE 
        sb.NAME = '油墩港(东大盈船闸)' 
        AND da.PASSTIME >= TO_DATE('2025-01-01 00:00:00', 'YYYY-MM-DD HH24:MI:SS')
        AND da.PASSTIME <= TO_DATE('2025-01-31 23:59:59', 'YYYY-MM-DD HH24:MI:SS');"""
        },
        {
            "input": "查询油墩港(东大盈船闸)2025年 1月份各类预警的案件总数与占比",
            "query": """WITH filtered_data AS (
      SELECT 
        zwt.NAME AS WARNING_TYPE_NAME
      FROM 
        BAYONET_DYNAMIC.DATAFUSION df
      JOIN 
        BAYONET_BUSSINESS.BOAT_WARNING bw 
        ON df.CODE = bw.EVENTCODE
      JOIN 
        BAYONET_BUSSINESS.ZD_WARNING_TYPE zwt  
        ON bw.WARNING_TYPE = zwt.ID  
      JOIN 
        BAYONET_BASICS.SYS_BAYONET sb  -- 关联卡口基础信息表
        ON df.BAYONET_ID = sb.ID       -- 通过ID关联
      WHERE 
        sb.NAME = '油墩港(东大盈船闸)'  -- 直接使用卡口名称过滤
        AND df.PASSTIME >= DATE '2025-01-01' 
        AND df.PASSTIME < DATE '2025-02-01'
    )

    SELECT  WARNING_TYPE_NAME,  COUNT(*) AS record_count,  ROUND(COUNT(*) * 100.0 / SUM(COUNT(*)) OVER (),2) AS percentage FROM filtered_data GROUP BY WARNING_TYPE_NAME ORDER BY  WARNING_TYPE_NAME;
    """
        }

    ]

    few_shot_example_prompt = PromptTemplate.from_template("输入问题: {input}\n生成SQL: {query}")
    few_shot_prompt = FewShotPromptTemplate(
        examples=examples,
        example_prompt=few_shot_example_prompt,
        suffix="\n---\n问题: {input}\n对应SQL: ",
        input_variables=["input", "table_info"],
    )

    few_shot_sql = few_shot_prompt.format(input=question, table_info=db.table_info)

    data = {
        "model": MODEL_NAME,
        "messages": [
            {"role": "system", "content": "你是一个Oracle SQL专家。请根据以下表结构生成对应问题的Oracle SQL语句，Oracle版本较低(11g及以下)，且尽量符合标准语法,不考虑ISDELETE和ISFINISH字段，如果没有强调年份，默认为2025年。\n\n以下是表结构（table_info）：\n{table_info}\n\n这是一些问题及其对应的SQL样例:"},
            {"role": "user", "content": few_shot_sql}
        ],
        "stream": True
    }

    response = requests.post(
        API_URL,
        data=json.dumps(data),
        stream=True,
        headers={'Content-Type': 'application/json'}
    )
    response.raise_for_status()

    # 流式处理并实时输出
    final_report = []
    print("\n开始生成sql：")
    for chunk in response.iter_lines():
        if chunk:
            decoded = json.loads(chunk.decode('utf-8'))
            if 'message' in decoded and 'content' in decoded['message']:
                content = decoded['message']['content']
                print(content, end="", flush=True)  # 实时输出
                final_report.append(content)

    complete_sql = ''.join(final_report).strip()
    print("\n数据分析完成")
    return complete_sql

# 使用示例
if __name__ == "__main__":
    question = "统计一月份各卡口点的船舶通过数量，按照通过船舶数量降序排列，只显示通过量前10的卡口。"

    sql = generate_sql_data(question)
    # print("\n最终sql内容：\n", sql)





    # if "```sql" in llm_response:
    #     extracted_sql = extract_sql_from_markdown(llm_response)
    #     print("提取出的 SQL:")
    #     print(extracted_sql)
    # else:
    #     extracted_sql = llm_response.replace(";", "")
    #
    # try:
    #     # 执行SQL
    #     sql_result = db.run(extracted_sql)
    #     print("SQL结果:", sql_result)
    #
    # except Exception as e:
    #     print(f"执行查询时出错: {e}")